package com.atguigu;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class UnionFind {
    //样本进来会包一层，叫做元素
    public static class Element<V>{
        public V value;
        public Element(V value){
            this.value=value;
        }
    }
    public static class UnionFindSet<V>{
        public HashMap<V,Element<V>> elementHashMap;
        //key 某个元素 value 该元素的父
        public HashMap<Element<V>,Element<V>> fatherMap;
        // key 某个集合的代表元素， value 该集合的大小
        public HashMap<Element<V>,Integer> sizeMap;
        public UnionFindSet(List<V> list){
            for(V value:list){
                Element<V> element=new Element<>(value);
                elementHashMap.put(value,element);
                fatherMap.put(element,element);
                sizeMap.put(element,1);
            }
        }
        public void union(V a,V b){
            if(elementHashMap.containsKey(a) && elementHashMap.containsKey(b)) {
                Element<V> headA = elementHashMap.get(a);
                Element<V> headB=elementHashMap.get(b);
                if(headA!=headB) {
                    Element<V> big = sizeMap.get(headA) > sizeMap.get(headB) ? headA : headB;
                    Element<V> small = big == headA ? headB : headA;
                    fatherMap.put(small, big);
                    sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
                    sizeMap.remove(small);
                }
            }
        }
        public boolean isSameSet(V a,V b){
            if(elementHashMap.containsKey(a) && elementHashMap.containsKey(b)){
                return getHead(elementHashMap.get(a))==getHead(elementHashMap.get(b));
            }
            return false;
        }
        public Element<V> getHead(Element<V> e){
            Stack<Element<V>> stack=new Stack<>();
            while(e!=fatherMap.get(e)){
                stack.push(e);
                e=fatherMap.get(e);
            }
            //此时e就是头顶节点
            while(!stack.isEmpty()){
                fatherMap.put(stack.pop(),e);
            }
            return e;
        }
    }
}
